import os
from pathlib import Path

import pytest

from storage.kvstore.pickle import PickleKVStore


def test_pickle_kv_store_get_put(kv_path):
    store = PickleKVStore(kv_path)
    with store.tx():
        store.put(b'key', b'value')
    with store.tx():
        assert store.get(b'key') == b'value'


def test_pickle_kv_store_getitem_setitem(kv_path):
    store = PickleKVStore(kv_path)
    with store.tx():
        store[b'key'] = b'value'
    with store.tx():
        assert store[b'key'] == b'value'


def test_pickle_kv_store_delete(kv_path):
    store = PickleKVStore(kv_path)
    with store.tx():
        store.put(b'key', b'value')

    with store.tx():
        del store[b'key']

    with store.tx():
        assert store.get(b'key') is None


@pytest.mark.benchmark
def test_benchmark_read_pickle(benchmark, insert_keys, get_keys):
    db_path = Path('./testdata.pickle.bin')
    if db_path.exists():
        db_path.unlink()

    db = PickleKVStore(db_path)
    with db.tx():
        for k in insert_keys:
            db[k] = os.urandom(8)

    def benchmark_pickle_get():
        nonlocal db
        with db.tx():
            for k in get_keys:
                _ = db.get(k)

    benchmark(benchmark_pickle_get)


@pytest.mark.benchmark
def test_benchmark_write_pickle(benchmark, insert_keys):
    db_path = Path('./testdata.pickle.bin')

    def benchmark_func():
        if db_path.exists():
            db_path.unlink()

        db = PickleKVStore(str(db_path))
        for k in insert_keys:
            db[k] = os.urandom(8)

    benchmark(benchmark_func)
